package com.example.plantuml;

import java.io.ByteArrayOutputStream;
import java.io.IOException;
import java.io.InputStream;
import java.net.HttpURLConnection;
import java.net.URI;
import java.net.URL;
import java.nio.charset.StandardCharsets;
import java.util.zip.Deflater;

/**
 * Minimal PlantUML server client that implements the standard deflate + base64-like encoding.
 * Reference: https://www.plantuml.com/plantuml
 */
public final class PlantUmlClient {
    private final String serverBaseUrl;
    private volatile long lastRequestAtMs = 0L;
    private static final int MAX_RETRIES = 5;
    private static final String USER_AGENT = "plantuml-docx-converter/1.0 (+https://example.com)";
    private final boolean localServer;

    public PlantUmlClient(String serverBaseUrl) {
        String base = serverBaseUrl;
        if (base.endsWith("/")) {
            base = base.substring(0, base.length() - 1);
        }
        this.serverBaseUrl = base;
        this.localServer = isLocalBase(base);
    }

    public byte[] renderPng(String umlSource) throws IOException {
        // Local PicoWeb only supports GET; for remote servers prefer POST then fallback to GET.
        if (localServer) {
            return renderViaGetWithRetries(umlSource);
        } else {
            try {
                return renderViaPostWithRetries(umlSource);
            } catch (IOException ignored) {
                return renderViaGetWithRetries(umlSource);
            }
        }
    }

    private byte[] renderViaPostWithRetries(String umlSource) throws IOException {
        String url = serverBaseUrl + "/png";
        byte[] bytes = umlSource.getBytes(StandardCharsets.UTF_8);
        IOException last = null;
        for (int attempt = 0; attempt <= MAX_RETRIES; attempt++) {
            throttle();
            try {
                HttpURLConnection conn = (HttpURLConnection) new URL(url).openConnection();
                conn.setRequestMethod("POST");
                conn.setDoOutput(true);
                conn.setConnectTimeout(15000);
                conn.setReadTimeout(30000);
                conn.setRequestProperty("User-Agent", USER_AGENT);
                conn.setRequestProperty("Accept", "image/png,*/*;q=0.8");
                conn.setRequestProperty("Cache-Control", "no-cache");
                conn.setRequestProperty("Pragma", "no-cache");
                conn.setRequestProperty("Connection", "close");
                conn.setRequestProperty("Content-Type", "text/plain; charset=UTF-8");
                conn.getOutputStream().write(bytes);
                int code = conn.getResponseCode();
                if (code >= 200 && code < 300) {
                    try (InputStream in = conn.getInputStream()) {
                        byte[] body = readAllBytes(in);
                        ensurePng(body);
                        return body;
                    }
                }
                if (shouldRetry(code) && attempt < MAX_RETRIES) {
                    sleepBackoff(attempt, code);
                    continue;
                }
                throw new IOException("HTTP " + code + " from server on POST");
            } catch (IOException e) {
                last = e;
                if (attempt < MAX_RETRIES) {
                    sleepBackoff(attempt, null);
                    continue;
                }
            }
        }
        throw last != null ? last : new IOException("POST failed with unknown error");
    }

    private byte[] renderViaGetWithRetries(String umlSource) throws IOException {
        String encoded = encodeForPlantUml(umlSource);
        String url = serverBaseUrl + "/png/" + encoded;
        IOException last = null;
        for (int attempt = 0; attempt <= MAX_RETRIES; attempt++) {
            throttle();
            try {
                HttpURLConnection conn = (HttpURLConnection) new URL(url).openConnection();
                conn.setRequestMethod("GET");
                conn.setConnectTimeout(15000);
                conn.setReadTimeout(30000);
                conn.setRequestProperty("User-Agent", USER_AGENT);
                conn.setRequestProperty("Accept", "image/png,*/*;q=0.8");
                conn.setRequestProperty("Cache-Control", "no-cache");
                conn.setRequestProperty("Pragma", "no-cache");
                conn.setRequestProperty("Connection", "close");
                int code = conn.getResponseCode();
                if (code >= 200 && code < 300) {
                    try (InputStream in = conn.getInputStream()) {
                        byte[] body = readAllBytes(in);
                        ensurePng(body);
                        return body;
                    }
                }
                if (shouldRetry(code) && attempt < MAX_RETRIES) {
                    sleepBackoff(attempt, code);
                    continue;
                }
                throw new IOException("HTTP " + code + " from server on GET");
            } catch (IOException e) {
                last = e;
                if (attempt < MAX_RETRIES) {
                    sleepBackoff(attempt, null);
                    continue;
                }
            }
        }
        throw last != null ? last : new IOException("GET failed with unknown error");
    }

    private void throttle() {
        if (localServer) {
            return; // local server: skip throttle for speed
        }
        long now = System.currentTimeMillis();
        long delta = now - lastRequestAtMs;
        int delay = java.util.concurrent.ThreadLocalRandom.current().nextInt(1000, 2001);
        if (delta < delay) {
            try { Thread.sleep(delay - delta); } catch (InterruptedException ignored) {}
        }
        lastRequestAtMs = System.currentTimeMillis();
    }

    private static boolean shouldRetry(int code) {
        return code == 520 || code == 429 || code == 503 || code == 502 || code == 504;
    }

    private static void sleepBackoff(int attempt, Integer statusCode) {
        long baseMs = (statusCode != null && statusCode == 520) ? 2000L : 1500L; // 2s for 520, 1.5s default
        long backoff = baseMs * (1L << Math.min(attempt, 3)); // 1x,2x,4x,8x
        int jitter = java.util.concurrent.ThreadLocalRandom.current().nextInt(-300, 901);
        long sleep = Math.max(1500L, backoff + jitter);
        try { Thread.sleep(sleep); } catch (InterruptedException ignored) {}
    }

    private static boolean isLocalBase(String baseUrl) {
        try {
            URI uri = URI.create(baseUrl);
            String host = uri.getHost();
            if (host == null) return false;
            return "localhost".equalsIgnoreCase(host) || "127.0.0.1".equals(host);
        } catch (Exception ignored) {
            return false;
        }
    }

    private static byte[] readAllBytes(InputStream in) throws IOException {
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        byte[] buf = new byte[8192];
        int r;
        while ((r = in.read(buf)) != -1) {
            bos.write(buf, 0, r);
        }
        return bos.toByteArray();
    }

    private static void ensurePng(byte[] bytes) throws IOException {
        if (bytes == null || bytes.length < 8) {
            throw new IOException("Empty response from server");
        }
        int[] sig = {0x89, 0x50, 0x4E, 0x47, 0x0D, 0x0A, 0x1A, 0x0A};
        for (int i = 0; i < sig.length; i++) {
            if ((bytes[i] & 0xFF) != sig[i]) {
                String head = new String(bytes, 0, Math.min(bytes.length, 64), StandardCharsets.ISO_8859_1)
                        .replaceAll("\n|\r", " ");
                throw new IOException("Server response is not PNG (starts with '" + head + "')");
            }
        }
    }

    // The following encoding logic is adapted from the public PlantUML specification.
    private static String encodeForPlantUml(String text) {
        byte[] data = text.getBytes(StandardCharsets.UTF_8);
        byte[] deflated = deflate(data);
        return encode64(deflated);
    }

    private static byte[] deflate(byte[] data) {
        Deflater deflater = new Deflater(Deflater.BEST_COMPRESSION, true);
        deflater.setInput(data);
        deflater.finish();
        byte[] buffer = new byte[1024];
        ByteArrayOutputStream bos = new ByteArrayOutputStream();
        while (!deflater.finished()) {
            int count = deflater.deflate(buffer);
            bos.write(buffer, 0, count);
        }
        deflater.end();
        return bos.toByteArray();
    }

    private static String encode64(byte[] data) {
        StringBuilder r = new StringBuilder();
        for (int i = 0; i < data.length; i += 3) {
            int b1 = data[i] & 0xFF;
            if (i + 2 == data.length) {
                int b2 = data[i + 1] & 0xFF;
                r.append(append3bytes(b1, b2, 0));
            } else if (i + 1 == data.length) {
                r.append(append3bytes(b1, 0, 0));
            } else {
                int b2 = data[i + 1] & 0xFF;
                int b3 = data[i + 2] & 0xFF;
                r.append(append3bytes(b1, b2, b3));
            }
        }
        return r.toString();
    }

    private static String append3bytes(int b1, int b2, int b3) {
        int c1 = b1 >> 2;
        int c2 = ((b1 & 0x3) << 4) | (b2 >> 4);
        int c3 = ((b2 & 0xF) << 2) | (b3 >> 6);
        int c4 = b3 & 0x3F;
        StringBuilder r = new StringBuilder();
        r.append(encode6bit(c1 & 0x3F));
        r.append(encode6bit(c2 & 0x3F));
        r.append(encode6bit(c3 & 0x3F));
        r.append(encode6bit(c4 & 0x3F));
        return r.toString();
    }

    private static char encode6bit(int b) {
        if (b < 10) return (char) ('0' + b);
        b -= 10;
        if (b < 26) return (char) ('A' + b);
        b -= 26;
        if (b < 26) return (char) ('a' + b);
        b -= 26;
        if (b == 0) return '-';
        if (b == 1) return '_';
        return '?';
    }
}


